#入口
import os
import threading
import time
from datetime import datetime
from PyQt5.QtWidgets import QApplication, QWidget, QButtonGroup, QMessageBox, QPlainTextEdit
from PyQt5.QtCore import *
from PyQt5.QtGui import QIcon, QRegExpValidator
from PyQt5 import uic
import sys
import re
from PyQt5.QtCore import QTimer
import S_DES
from S_DES import *
from test import *

"""
PyQt 多线程
"""
class WorkerSignals(QObject):
    # progress: (worker_id:int, percent:int)
    progress = pyqtSignal(int, int)
    # found: (worker_id:int, key_str:str)
    found = pyqtSignal(int, str)
    finished = pyqtSignal(int)
    log = pyqtSignal(str)

class Worker(QObject):
    """
    一个可被移动到 QThread 的工作对象。每个 Worker 枚举一段 key 空间并尝试对多个 (plain,cipher) 对进行短路验证。
    """
    def __init__(self, worker_id, key_start, key_end, pairs, stop_event):
        super().__init__()
        self.worker_id = worker_id
        self.key_start = key_start  # inclusive
        self.key_end = key_end      # exclusive
        # pairs: list of tuples [(plain_bin_str, cipher_bin_str), ...]
        self.pairs = pairs
        self.signals = WorkerSignals()
        self.stop_event = stop_event

    @pyqtSlot()
    def run(self):
        """长循环在这里执行，必须响应 stop_event 以便能中止。"""
        total = max(1, self.key_end - self.key_start)
        checked = 0
        # 控制进度更新频率：最多 ~200 次更新
        report_interval = max(1, total // 200)

        try:
            for k in range(self.key_start, self.key_end):
                if self.stop_event.is_set():
                    self.signals.log.emit(f"Worker {self.worker_id}: 停止标志检测到，退出。")
                    break

                key_bin = format(k, '010b')   # 10 位密钥字符串
                # 每个 key 只 new 一次 DES 实例（节省 round-key 生成开销）
                try:
                    DES = S_DES_Binary(key_bin)
                except Exception as e:
                    # 如果 S_DES 初始化失败，记录并跳过该 key
                    self.signals.log.emit(f"Worker {self.worker_id}: 初始化 S_DES 出错 (key={key_bin})：{e}")
                    checked += 1
                    if (checked % report_interval) == 0 or checked == total:
                        percent = int(checked * 100 / total)
                        self.signals.progress.emit(self.worker_id, percent)
                    continue

                valid = True
                # 对所有 pairs 短路验证：一旦有不匹配就丢弃该 key
                for plain_bin, cipher_bin in self.pairs:
                    DES.Encryption(plain_bin)
                    if DES.ciperText != cipher_bin:
                        valid = False
                        break

                checked += 1
                if valid:
                    # 找到密钥，通知主线程
                    self.signals.found.emit(self.worker_id, key_bin)
                    #self.stop_event.set()
                    self.signals.log.emit(f"Worker {self.worker_id}: 找到密钥 {key_bin}")


                # 节流发送进度
                if (checked % report_interval) == 0 or checked == total:
                    percent = int(checked * 100 / total)
                    self.signals.progress.emit(self.worker_id, percent)

            # 最后把进度推到 100%（如果循环完整结束）
            self.signals.progress.emit(self.worker_id, 100)
        except Exception as e:
            self.signals.log.emit(f"Worker {self.worker_id} 异常: {e}")
        finally:
            self.signals.finished.emit(self.worker_id)

def attach_plaintext_limiter(widget: QPlainTextEdit, max_lines=8, max_len=8):
    """
    限制 QPlainTextEdit 输入：
      - 只允许 0/1
      - 每行最多 max_len 个字符
      - 最多 max_lines 行
      - 可正常回车换行
    """
    def on_text_changed():
        cursor = widget.textCursor()
        text = widget.toPlainText()

        # 拆分行
        lines = text.splitlines()
        modified = False

        # 限制行数
        if len(lines) > max_lines:
            lines = lines[:max_lines]
            modified = True

        new_lines = []
        for line in lines:
            filtered = re.sub(r'[^01]', '', line)
            if len(filtered) > max_len:
                filtered = filtered[:max_len]
            if filtered != line:
                modified = True
            new_lines.append(filtered)

        if modified:
            block = cursor.blockNumber()
            col = cursor.positionInBlock()

            widget.blockSignals(True)
            widget.setPlainText("\n".join(new_lines))
            widget.blockSignals(False)

            # 恢复光标位置
            doc = widget.document()
            blocks = min(block, len(new_lines) - 1)
            new_block = doc.findBlockByNumber(blocks)
            new_cursor = widget.textCursor()
            new_cursor.setPosition(new_block.position() + min(col, len(new_lines[blocks])))
            widget.setTextCursor(new_cursor)

    widget.textChanged.connect(on_text_changed)

class AnalysisWorker(QObject):
    # 在子线程运行封闭性分析防止堵塞主线程UI
    log = pyqtSignal(str)
    finished = pyqtSignal()

    def __init__(self, plain0):
        super().__init__()
        self.plain0 = plain0
        self._stop = False

    @pyqtSlot()
    def run(self):
        try:
            self.log.emit(f"=== 1) 固定明文统计: {self.plain0}")
            cnt, vmin, vmax, vavg, distinct = stats_for_plain(self.plain0)
            self.log.emit(f"产生的不同密文数量: {distinct}")
            self.log.emit(f"最小计数={vmin}, 最大计数={vmax}, 平均={vavg}\n")

            # Step 2
            self.log.emit("=== 2) 构造全部映射并分组等价密钥 ...")
            t0 = time.perf_counter()
            mapping_to_keys = build_full_mappings(save_progress=True,log = self.log)
            t1 = time.perf_counter()
            self.log.emit(f"完成，用时 {t1 - t0:.2f} 秒。不同映射数量: {len(mapping_to_keys)}\n")

            # Step 3
            self.log.emit("=== 3) 分析等价类 ...")
            sizes, diffs = analyze_equivalence_classes(mapping_to_keys)
            self.log.emit(f"等价类大小分布: {dict(sizes)}")
            if len(diffs) == 1:
                self.log.emit(f"所有大小为2的密钥对均相差恒定 XOR = {next(iter(diffs))}\n")
            else:
                self.log.emit(f"不同 XOR 值数量: {len(diffs)}\n")

            # Step 4
            self.log.emit("=== 4) Monte-Carlo 试验中（可能需要较长时间） ...")
            mc_t_values = (1, 2, 3, 4, 5, 6, 7)
            mc_stats = monte_carlo_trial(trials=500, t_values=mc_t_values,log = self.log)
            for t in mc_t_values:
                s = mc_stats[t]
                self.log.emit(f"t={t}: mean={s['mean']:.3f}, median={s['median']}, max={s['max']}\n")

            # Step 5
            self.log.emit("=== 5) 示例等价类 ===")
            shown = 0
            for mapping, keys in mapping_to_keys.items():
                if len(keys) > 1:
                    self.log.emit(f"等价类大小={len(keys)} 密钥={keys[:10]}")
                    shown += 1
                    if shown >= 10:
                        break

            self.log.emit("分析完成。")

        except Exception as e:
            self.log.emit(f"分析异常: {e}")
        finally:
            self.finished.emit()
class MainWindow(QWidget):
    def __init__(self):
        super().__init__()
        self.init_ui()
        self.switch_input_mode()
        # 线程管理相关
        self.threads = []  # list of QThread
        self.workers = []  # list of Worker
        self.worker_progress = {}  # worker_id -> percent
        self.stop_event = threading.Event()
        self.found_keys = []  # 保存找到的密钥（理论上最多几个）
        self.setWindowTitle("S-DES演示")
        self.tabWidget.setCurrentIndex(0)


    def init_ui(self):
        self.setWindowIcon(QIcon('./ui/keys.png'))
        self.ui = uic.loadUi('./ui/mainWindow.ui', self)
        #加解密绑定
        self.Binary = self.ui.radioButton
        self.Ascii = self.ui.radioButton_2
        self.Encryption = self.ui.radioButton_3
        self.Decryption = self.ui.radioButton_4
        self.input_group = QButtonGroup(self)
        self.input_group.addButton(self.Binary)
        self.input_group.addButton(self.Ascii)
        self.Binary.setChecked(True)
        self.mode_group = QButtonGroup(self)
        self.mode_group.addButton(self.Encryption)
        self.mode_group.addButton(self.Decryption)
        self.Encryption.setChecked(True)
        self.text = self.ui.lineEdit
        self.key = self.ui.lineEdit_2
        self.result = self.ui.textBrowser
        self.start = self.ui.pushButton

        #暴力破解页面绑定
        self.plainText_power = self.ui.plainTextEdit
        self.cipherText_power = self.ui.plainTextEdit_2
        attach_plaintext_limiter(self.plainText_power, max_lines=1024, max_len=8)
        attach_plaintext_limiter(self.cipherText_power, max_lines=1024, max_len=8)
        self.result_key = self.ui.textBrowser_3
        self.btnPower = self.ui.pushButton_2
        self.progressBar = self.ui.progressBar

        #封闭性测试的页面绑定
        self.plainTextFixed = self.ui.lineEdit_3
        self.btn = self.ui.pushButton_3
        self.analysis = self.ui.textBrowser_2

        # 事件绑定
        self.start.clicked.connect(self.startSDES)#开始
        self.Ascii.toggled.connect(self.switch_input_mode)#切换模式
        self.Binary.toggled.connect(self.switch_input_mode)#切换模式
        self.btnPower.clicked.connect(self.startPower)#开始暴力破解
        self.btn.clicked.connect(self.startAnalyse)
        # 补充：key的正则检验
        key_binary_validator = QRegExpValidator(QRegExp("[01]{0,10}")) # 0/1.最多10位
        self.key.clear()
        self.key.setValidator(key_binary_validator)
        # 补充：封闭性分析的正则检验
        test_binary_validator = QRegExpValidator(QRegExp("[01]{0,8}"))  # 0/1.最多10位
        self.plainTextFixed.setValidator(test_binary_validator)



    def startSDES(self):
        self.result.clear()
        text_CP = self.text.text()
        key = self.key.text()
        result = ""
        if len(key) != 10:
            QMessageBox.information(self,"提示","秘钥需10位,请重新输入", QMessageBox.Ok)
            return
        if self.Binary.isChecked():
            if len(text_CP) != 8 :
                QMessageBox.information(self,"提示","秘钥需8位,请重新输入", QMessageBox.Ok)
                return
            DES = S_DES_Binary(key)
            if self.Encryption.isChecked():
                DES.Encryption(text_CP)
                result = DES.ciperText
            elif self.Decryption.isChecked():
                DES.Decryption(text_CP)
                result = DES.plainText
        elif self.Ascii.isChecked():
            if self.Encryption.isChecked():
                result = S_DES.AsciiEncryption(text_CP,key)
            elif self.Decryption.isChecked():
                result = S_DES.AsciiDencryption(text_CP,key)
        self.showResult(text_CP,result)
    def switch_input_mode(self):
        # 设置验证器
        ascii_validator = None
        binary_validator = QRegExpValidator(QRegExp("[01]{0,8}"))  # 0/1，最多8位
        if self.Ascii.isChecked():
            self.text.clear()
            self.text.setValidator(ascii_validator)
        elif self.Binary.isChecked():
            self.text.clear()
            self.text.setValidator(binary_validator)
        self.result.clear()
    def showResult(self,text,result):
        self.result.clear()
        ans = ""
        if self.Decryption.isChecked():
            ans = text + "的解密结果是:\n"+result
        elif self.Encryption.isChecked():
            ans = text + "的加密结果是:\n" + result
        self.result.setText(ans)

    def startPower(self):
        """在点击暴力破解按钮时调用：开始或取消破解（支持 1..n 个明密文对）"""
        # 如果当前已有在跑的线程，则当作取消请求
        if any(t.isRunning() for t in self.threads):
            # 发出取消信号
            self.stop_event.set()
            self.result_key.append("已请求取消，等待线程退出...\n")
            self.btnPower.setEnabled(False)  # 直到线程真正结束
            return

        # 否则开始新的暴力破解任务
        plain_input = self.plainText_power.toPlainText().strip()
        cipher_input = self.cipherText_power.toPlainText().strip()


        def parse_pairs(plain_str, cipher_str):
            """
            解析两个 QTextEdit 输入的多行明文与密文，逐行配对
            返回 list[tuple[str, str]]
            """
            def clean(s: str) -> str:
                if s is None:
                    return ""
                s = s.replace('\ufeff', '').replace('\u200b', '')
                s = s.replace('\u3000', ' ')  # 全角空格
                s = ''.join(ch for ch in s if ch == '\n' or ord(ch) >= 32)
                return s.strip()

            plain_str = clean(plain_str)
            cipher_str = clean(cipher_str)

            # 拆分行，过滤空行
            plain_lines = [ln.strip() for ln in plain_str.splitlines() if ln.strip()]
            cipher_lines = [ln.strip() for ln in cipher_str.splitlines() if ln.strip()]

            if not plain_lines or not cipher_lines:
                raise ValueError("明文或密文输入为空。请至少输入一行 8 位二进制字符串。")

            if len(plain_lines) != len(cipher_lines):
                raise ValueError(f"明文和密文行数不一致：明文 {len(plain_lines)} 行，密文 {len(cipher_lines)} 行。")

            pairs = []
            for i, (p, c) in enumerate(zip(plain_lines, cipher_lines), start=1):
                if len(p) != 8 or len(c) != 8:
                    raise ValueError(f"第 {i} 行：明文或密文长度不是 8 位。\n明文={p!r}，密文={c!r}")
                if any(ch not in '01' for ch in p + c):
                    raise ValueError(f"第 {i} 行：明文或密文含有非二进制字符。\n明文={p!r}，密文={c!r}")
                pairs.append((p, c))

            return pairs

        try:
            pairs = parse_pairs(plain_input, cipher_input)
        except ValueError as e:
            QMessageBox.information(self, "输入错误", str(e), QMessageBox.Ok)
            return

        # 基本校验：每对都应该是 8 位二进制字符串
        for p, c in pairs:
            if len(p) != 8 or len(c) != 8 or any(ch not in '01' for ch in p + c):
                QMessageBox.information(self, "提示", "明文/密文需为 8 位二进制（每对）", QMessageBox.Ok)
                return

        # reset 状态
        key_space = 1024
        self.progressBar.setValue(0)
        self.result_key.clear()
        self.found_keys = []
        self.worker_progress.clear()
        self.stop_event.clear()
        # 启用取消行为（按钮文本切换）
        self.btnPower.setText("取消")
        self.btnPower.setEnabled(True)
        self.start_time_perf = time.perf_counter()
        self.start_time_clock = datetime.now()
        start_ts = self.start_time_clock.strftime("%Y-%m-%d %H:%M:%S.%f")[:-3]
        self.result_key.append(f"[{start_ts}+0ms] 开始暴力破解(keyspace={key_space},线程数={4})")
        # 划分 key 空间（0..1023）
        num_workers = 4
        chunk = (key_space + num_workers - 1) // num_workers

        # 清理旧线程容器（若之前线程都结束则没有问题）
        self.threads = []
        self.workers = []

        for i in range(num_workers):
            start = i * chunk
            end = min(start + chunk, key_space)
            if start >= end:
                continue
            worker = Worker(i, start, end, pairs, self.stop_event)
            thread = QThread()
            worker.moveToThread(thread)

            # 连接信号
            worker.signals.progress.connect(self.handle_worker_progress)
            worker.signals.found.connect(self.handle_worker_found)
            worker.signals.finished.connect(self.handle_worker_finished)
            worker.signals.log.connect(self.handle_worker_log)

            thread.started.connect(worker.run)
            # 当线程结束时自动退出并释放
            worker.signals.finished.connect(thread.quit)

            self.threads.append(thread)
            self.workers.append(worker)

            # 初始化该 worker 的进度
            self.worker_progress[i] = 0

        # 启动所有线程
        for th in self.threads:
            th.start()

    def handle_worker_progress(self, worker_id, percent):
        # 更新单个 worker 的进度，然后更新总进度（平均法）
        self.worker_progress[worker_id] = percent
        # 计算平均进度（只按参与的 worker）
        active_ids = list(self.worker_progress.keys())
        total_percent = sum(self.worker_progress[i] for i in active_ids) / max(1, len(active_ids))
        self.progressBar.setValue(int(total_percent))

    def handle_worker_found(self, worker_id, key_str):
        # 收到某个线程找到密钥
        now = datetime.now()
        now_str = now.strftime('%Y-%m-%d %H:%M:%S.%f')[:-3]
        elapsed_ms = int((time.perf_counter() - getattr(self, 'start_time_perf', time.perf_counter())) * 1000)
        if key_str not in self.found_keys:
            self.found_keys.append(key_str)
            # 写入带时间戳的发现信息
            self.result_key.append(f"[{now_str} +{elapsed_ms}ms] 找到密钥（worker {worker_id}）: {key_str}\n")
        else:
            # 若重复发现也记录（可选）
            self.result_key.append(f"[{now_str} +{elapsed_ms}ms] 重复发现密钥（worker {worker_id}）: {key_str}\n")
        # 切换按钮文本，允许用户等待线程退出或直接禁用
        self.btnPower.setText("开始暴力破解")
        self.btnPower.setEnabled(True)

    def handle_worker_log(self, text):
        # 在主线程将 worker 发过来的日志加上时间戳和相对耗时（ms）
        now = datetime.now()
        now_str = now.strftime('%Y-%m-%d %H:%M:%S.%f')[:-3]
        elapsed_ms = int((time.perf_counter() - getattr(self, 'start_time_perf', time.perf_counter())) * 1000)
        self.result_key.append(f"[{now_str} +{elapsed_ms}ms] {text}\n")

    def handle_worker_finished(self, worker_id):
        # 当一个 worker 完成时检查是否所有线程都已结束
        still_running = any(t.isRunning() for t in self.threads)
        if not still_running:
            # 所有线程已结束：汇总结果、重置按钮状态，记录结束时间戳与总耗时
            now = datetime.now()
            now_str = now.strftime('%Y-%m-%d %H:%M:%S.%f')[:-3]
            total_elapsed_ms = int((time.perf_counter() - getattr(self, 'start_time_perf', time.perf_counter())) * 1000)

            if not self.found_keys:
                self.result_key.append(f"[{now_str} +{total_elapsed_ms}ms] 破解结束：未找到匹配密钥。\n")
            else:
                self.result_key.append(
                    f"[{now_str} +{total_elapsed_ms}ms] 破解结束：已找到密钥（共 {len(self.found_keys)} 个）。\n")
                self.result_key.append("候选密钥列表: " + ", ".join(self.found_keys) + "\n")

            self.btnPower.setText("开始暴力破解")
            self.btnPower.setEnabled(True)
            # 清除 stop_event 以便下一次能重新开始
            self.stop_event.clear()

    def append_analysis_text(self, text):
        """供 AnalysisWorker.log 信号调用的槽函数"""
        timestamp = datetime.now().strftime('%H:%M:%S')
        self.analysis.append(f"[{timestamp}] {text}")
        # 自动滚动到底部
        self.analysis.verticalScrollBar().setValue(self.analysis.verticalScrollBar().maximum())
    def startAnalyse(self):
        plain0 = self.plainTextFixed.text().strip()
        if not plain0:
            QMessageBox.information(self, "提示", "请输入明文", QMessageBox.Ok)
            return

        self.analysis.clear()
        self.analysis.append(f"[开始分析] 固定明文: {plain0}\n")

        # 禁用按钮，防止重复启动
        self.btn.setEnabled(False)

        # 创建线程与 worker
        self.analysis_thread = QThread()
        self.analysis_worker = AnalysisWorker(plain0)
        self.analysis_worker.moveToThread(self.analysis_thread)

        # 信号连接
        self.analysis_worker.log.connect(self.append_analysis_text)
        self.analysis_worker.finished.connect(self.analysis_thread.quit)
        self.analysis_worker.finished.connect(lambda: self.btn.setEnabled(True))
        self.analysis_worker.finished.connect(lambda: self.analysis.append("\n[分析结束]\n"))
        self.analysis_thread.started.connect(self.analysis_worker.run)

        # 线程结束后自动清理
        self.analysis_thread.finished.connect(self.analysis_thread.deleteLater)
        self.analysis_thread.start()
if __name__=="__main__":
    app = QApplication(sys.argv)
    window = MainWindow()
    window.show()
    sys.exit(app.exec_())